import os

import numpy as np
import skimage.io
from brisque_feature import brisque_feature
import scipy.io
from tqdm import trange
from sklearn.model_selection import GridSearchCV,cross_validate,ShuffleSplit
from sklearn.svm import SVR
from scipy import stats
import pandas as pd

def adjust_name(idx,ntype):
    if ntype=='jp2k':
        if idx>227:
            raise ValueError("Image Index Out of Range")
        return idx-1
    elif ntype=='jpeg':
        if idx>233:
            raise ValueError("Image Index Out of Range")
        return idx+226
    elif ntype=='wn':
        if idx>174:
            raise ValueError("Image Index Out of Range")
        return idx+459
    elif ntype=='gblur':
        if idx>174:
            raise ValueError("Image Index Out of Range")
        return idx+633
    elif ntype=='fastfading':
        if idx>174:
            raise ValueError("Image Index Out of Range")
        return idx+807
    else:
        raise ValueError("Unkown Distortion Type")

def CreatSet(image_folder, output_feat_path, output_score_path, excel_path):
    X = []
    Y = []
    # 读取Excel文件
    df = pd.read_excel(excel_path)

    for img_filename in os.listdir(image_folder):
        if img_filename.endswith(".jpg"):
            img_path = os.path.join(image_folder, img_filename)
            img = skimage.io.imread(img_path, as_gray=True)
            feat = brisque_feature(img, device='cpu')
            X.append(feat)
            # 假设Excel文件的第二列包含文件名
            # 检查文件名是否在Excel的第二列中
            matching_row = df[df.iloc[:, 1] == img_filename]
            if not matching_row.empty:
                # 获取第十五列的值
                excel_value = matching_row.iloc[0, 14]  # 这里假设第十五列索引为14
                Y.append(excel_value)
            else:
                # 如果没有找到匹配的行，可以选择添加一个默认值或抛出异常
                Y.append(np.nan)  # 或者抛出异常
    np.save(output_feat_path, np.array(X))
    np.save(output_score_path, np.array(Y))

def evaluate():
    x=np.load('feat_Live.npy')
    y=np.load('score_Live.npy')
    x_val = np.load('val_feat_Live.npy')
    y_val = np.load('val_score_Live.npy')
    ss = ShuffleSplit(n_splits=1000, random_state=0, test_size=0.2)
    cout_num = 0
    predictions = []
    params = [
        {'kernel':['rbf'], 'C': 2 ** (np.arange(-8, 8, 0.8)),'gamma': 2 ** (np.arange(-8, 8, 0.8))}
    ]
    svr = SVR()
    clf = GridSearchCV(svr, params)
    clf.fit(x, y)
    SROCC_box=np.zeros((1000,1),dtype=np.float32)
    PLCC_box=np.zeros((1000,1),dtype=np.float32)

    '''for train_index,test_index in ss.split(x):
        svr = SVR(kernel='rbf', C=clf.best_params_['C'], gamma=clf.best_params_['gamma'])
        svr.fit(x[train_index],y[train_index])
        predict=svr.predict(x[test_index])
        predictions.extend(predict)
        SROCC,_=stats.spearmanr(predict, y[test_index])
        PLCC,_ = stats.pearsonr(predict, y[test_index])
        SROCC_box[cout_num,:]=SROCC
        PLCC_box[cout_num,:]=PLCC
        cout_num=cout_num+1
        
    print('Median SRCC %.4f  MedianPLCC %.4f'%(np.median(SROCC_box),np.median(PLCC_box)))
    np.save('predictions_live.npy', np.array(predictions))'''

    for train_index, test_index in ss.split(x):
        svr = SVR(kernel='rbf', C=clf.best_params_['C'], gamma=clf.best_params_['gamma'])
        svr.fit(x[train_index], y[train_index])
        predict = svr.predict(x[test_index])
        predictions.extend(predict)
        SROCC, _ = stats.spearmanr(predict, y[test_index])
        PLCC, _ = stats.pearsonr(predict, y[test_index])
        SROCC_box[cout_num, :] = SROCC
        PLCC_box[cout_num, :] = PLCC
        cout_num = cout_num + 1
    predict = svr.predict(x_val)
    SROCC, _ = stats.spearmanr(predict, y_val)
    PLCC, _ = stats.pearsonr(predict, y_val)
    print('Median SRCC %.4f  MedianPLCC %.4f' % (np.median(SROCC_box), np.median(PLCC_box)))
    print('VAL SRCC %.4f  VAL PLCC %.4f' % (SROCC, PLCC))
    np.save('predictions_live.npy', np.array(predict))




if __name__=='__main__':
    local_image_folder = 'C:/Users/zlsjNKJS/Desktop/090/train'
    val_image_folder = 'C:/Users/zlsjNKJS/Desktop/090/val'

    #训练集
    CreatSet(local_image_folder, output_feat_path='feat_Live.npy', output_score_path='score_Live.npy', excel_path = "C:/Users/zlsjNKJS/Desktop/090/train.xlsx")
    '''data = np.load('score_Live.npy')
    df = pd.DataFrame(data)
    df.to_csv('C:/Users/zlsjNKJS/Desktop/090/train/output.csv', index=False)'''
    #验证集
    CreatSet(val_image_folder, output_feat_path='val_feat_Live.npy', output_score_path='val_score_Live.npy',
             excel_path="C:/Users/zlsjNKJS/Desktop/090/val.xlsx")
    evaluate()
    predictions = np.load('predictions_live.npy')
    for idx, pred in enumerate(predictions):
        print(f"Image {idx} predicted quality score:{pred:.4f}")
    data = np.load('predictions_live.npy')
    df = pd.DataFrame(data)
    df.to_csv('C:/Users/zlsjNKJS/Desktop/090/brisque.csv', index=False)